package com.fastexport.util;

import com.alibaba.excel.EasyExcel;
import com.alibaba.excel.ExcelWriter;
import com.alibaba.excel.write.builder.ExcelWriterSheetBuilder;
import com.alibaba.excel.write.metadata.WriteSheet;
import com.fastexport.constant.Watermark;
import com.fastexport.handler.PresetWatermarkSheetHandler;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import javax.servlet.http.HttpServletResponse;
import java.io.*;
import java.time.LocalDateTime;
import java.time.format.DateTimeFormatter;
import java.util.*;

public class FastExportUtils {

    private static final Logger logger = LoggerFactory.getLogger(FastExportUtils.class);

    public static void export(List<?> list) {
        HttpServletResponse response = HttpContextUtils.getHttpServletResponse();
        export(response, list);
    }

    public static void export(List<?> list, String fileName) {
        HttpServletResponse response = HttpContextUtils.getHttpServletResponse();
        export(response, list, fileName);
    }

    public static void export(HttpServletResponse response, List<?> list) {
        String fileName = LocalDateTime.now().format(DateTimeFormatter.ofPattern("yyyyMMddHHmmss"));
        export(response, list, fileName);
    }

    public static void export(HttpServletResponse response, List<?> list, String fileName) {
        String sheetName = "sheet1";
        export(response, list, fileName, sheetName);
    }

    public static void export(HttpServletResponse response, List<?> list, String fileName, String sheetName) {
        ExcelWriterSheetBuilder builder = getWriterSheetBuilder(response, list, fileName, sheetName);
        builder.doWrite(list);
    }

    public static ExcelWriterSheetBuilder getWriterSheetBuilder(HttpServletResponse response, List<?> list, String fileName, String sheetName) {
        ValidateUtils.validateInputs(response, list);
        fileName = ValidateUtils.fileNameHandle(fileName);
        BaseHandlerUtils.setResponse(response, fileName);

        Class<?> aClass = list.get(0).getClass();
        List<?> headList = BaseHandlerUtils.getHeadList(list);

        try (ByteArrayOutputStream bos = new ByteArrayOutputStream()) {
            // 写入内存流
            EasyExcel.write(bos, aClass)
                    .inMemory(true)
                    .registerWriteHandler(new PresetWatermarkSheetHandler(new Watermark()))
                    .sheet(sheetName)
                    .doWrite(headList);

            // 将内存流转换为InputStream
            InputStream templateInputStream = new ByteArrayInputStream(bos.toByteArray());

            // 直接从内存流读取模板并写入响应
            return EasyExcel.write(response.getOutputStream(), aClass)
                    .withTemplate(templateInputStream)
                    .sheet(sheetName)
                    .needHead(false);
        } catch (Exception e) {
            logger.error("Error during Excel export.", e);
            throw new RuntimeException("Error during Excel export.", e);
        }
    }

    public static void exportMultiple(HttpServletResponse response, Map<String, List<?>> map, String fileName) {
        ValidateUtils.validateInputs(response, map);
        fileName = ValidateUtils.fileNameHandle(fileName);
        BaseHandlerUtils.setResponse(response, fileName);

        try (ByteArrayOutputStream bos = new ByteArrayOutputStream()) {
            ExcelWriter excelBuffer = EasyExcel.write(bos, Object.class)
                    .inMemory(true)
                    .registerWriteHandler(new PresetWatermarkSheetHandler(new Watermark()))
                    .build();
            for (Map.Entry<String, List<?>> entry : map.entrySet()) {
                String sheetName = entry.getKey();
                List<?> dataList = entry.getValue();
                List<?> headList = BaseHandlerUtils.getHeadList(dataList);

                Class<?> aClass = headList.get(0).getClass();
                WriteSheet writeSheet = EasyExcel.writerSheet(sheetName).head(aClass).build();
                excelBuffer.write(headList, writeSheet);
            }
            excelBuffer.finish();

            // 将内存流转换为InputStream
            InputStream templateInputStream = new ByteArrayInputStream(bos.toByteArray());

            // 直接从内存流读取模板并写入响应
            ExcelWriter excelWriter = EasyExcel.write(response.getOutputStream())
                    .withTemplate(templateInputStream)
                    .registerWriteHandler(new PresetWatermarkSheetHandler(new Watermark()))
                    .build();

            for (Map.Entry<String, List<?>> entry : map.entrySet()) {
                String sheetName = entry.getKey();
                List<?> dataList = entry.getValue();

                WriteSheet writeSheet = EasyExcel.writerSheet(sheetName)
                        .needHead(false).build();
                excelWriter.write(dataList, writeSheet);
            }

            excelWriter.finish();
        } catch (Exception e) {
            logger.error("Error during Excel export.", e);
            throw new RuntimeException("Error during Excel export.", e);
        }
    }

}
